import numpy as np
import random
import networkx as nx


class RandomWalker(object):
    """
    随机游走的封装类，包含node2vec、deepwalk、 node2vec_kk 三种随机游走方案
    注意：deepwalk算法在这里并不可以使用weight权重，但是你可以使用node2vec的方案替代
    """

    def __init__(self, nx_G: nx.Graph, is_directed, p, q, walk_type) -> None:
        self.G = nx_G
        self.is_directed = is_directed
        self.p = p
        self.q = q
        if walk_type == 'node2vec':
            self.walk = Node2Vec(nx_G, is_directed, p, q, use_rejection_sampling=False)
            self.walk.preprocess_transition_probs()
        elif walk_type == 'deep_walk':
            self.walk = DeepWalk(nx_G, is_directed, p, q)
        elif walk_type == 'node2vec_kk':
            self.walk = Node2Vec(nx_G, is_directed, p, q, use_rejection_sampling=True)
            self.walk.preprocess_transition_probs()
        else:
            raise Exception("Need walk type settings")

    def simulate_walks(self, num_walks, walk_length):
        """
        模拟随机游走
        """
        G = self.G
        walks = []
        nodes = list(G.nodes())
        print('Walk iteration:')
        for walk_iter in range(num_walks):
            print(str(walk_iter + 1), '/', str(num_walks))
            random.shuffle(nodes)
            for node in nodes:
                walks.append(self.walk.walk(walk_length=walk_length, start_node=node))
        return walks


class DeepWalk(object):
    """
    深度优先随机游走
    """

    def __init__(self, nx_G: nx.Graph, is_directed, p, q) -> None:
        self.G = nx_G
        self.is_directed = is_directed
        self.p = p
        self.q = q

    def walk(self, walk_length, start_node):

        walk = [start_node]

        while len(walk) < walk_length:
            cur = walk[-1]
            cur_nbrs = list(self.G.neighbors(cur))
            if len(cur_nbrs) > 0:
                walk.append(random.choice(cur_nbrs))
            else:
                break
        return walk


class Node2Vec(object):
    """
    Node2vec的游走策略
    """

    def __init__(self, nx_G: nx.Graph, is_directed, p, q, use_rejection_sampling=False):
        self.G = nx_G
        self.is_directed = is_directed
        self.p = p
        self.q = q
        self.use_rejection_sampling = use_rejection_sampling

    def walk(self, walk_length, start_node):
        """
        封装两种策略
        :param walk_length:
        :param start_node:
        :return:
        """
        if not self.use_rejection_sampling:
            return self.walk1(walk_length, start_node)
        else:
            return self.walk2(walk_length, start_node)

    def walk2(self, walk_length, start_node):
        def rejection_sample(inv_p, inv_q, nbrs_num):
            upper_bound = max(1.0, max(inv_p, inv_q))
            lower_bound = min(1.0, min(inv_p, inv_q))
            shatter = 0
            second_upper_bound = max(1.0, inv_q)
            if inv_p > second_upper_bound:
                shatter = second_upper_bound / nbrs_num
                upper_bound = second_upper_bound + shatter
            return upper_bound, lower_bound, shatter

        G = self.G
        alias_nodes = self.alias_nodes
        inv_p = 1.0 / self.p
        inv_q = 1.0 / self.q
        walk = [start_node]
        while len(walk) < walk_length:
            cur = walk[-1]
            cur_nbrs = list(G.neighbors(cur))
            if len(cur_nbrs) > 0:
                if len(walk) == 1:
                    walk.append(
                        cur_nbrs[alias_draw(alias_nodes[cur][0], alias_nodes[cur][1])])
                else:
                    upper_bound, lower_bound, shatter = rejection_sample(
                        inv_p, inv_q, len(cur_nbrs))
                    prev = walk[-2]
                    prev_nbrs = set(G.neighbors(prev))
                    while True:
                        prob = random.random() * upper_bound
                        if prob + shatter >= upper_bound:
                            next_node = prev
                            break
                        next_node = cur_nbrs[alias_draw(
                            alias_nodes[cur][0], alias_nodes[cur][1])]
                        if prob < lower_bound:
                            break
                        if prob < inv_p and next_node == prev:
                            break
                        _prob = 1.0 if next_node in prev_nbrs else inv_q
                        if prob < _prob:
                            break
                    walk.append(next_node)
            else:
                break
        return walk

    def walk1(self, walk_length, start_node):
        """
        Simulate a random walk starting from start node.
        """
        G = self.G
        alias_nodes = self.alias_nodes
        alias_edges = self.alias_edges

        walk = [start_node]

        while len(walk) < walk_length:
            cur = walk[-1]
            cur_nbrs = sorted(G.neighbors(cur))
            if len(cur_nbrs) > 0:
                if len(walk) == 1:
                    walk.append(cur_nbrs[alias_draw(alias_nodes[cur][0], alias_nodes[cur][1])])
                else:
                    prev = walk[-2]
                    next = cur_nbrs[alias_draw(alias_edges[(prev, cur)][0],
                                               alias_edges[(prev, cur)][1])]
                    walk.append(next)
            else:
                break

        return walk

    def get_alias_edge(self, src, dst):
        """
        Get the alias edge setup lists for a given edge.
        """
        G = self.G
        p = self.p
        q = self.q

        unnormalized_probs = []
        for dst_nbr in sorted(G.neighbors(dst)):
            if dst_nbr == src:
                unnormalized_probs.append(G[dst][dst_nbr]['weight'] / p)
            elif G.has_edge(dst_nbr, src):
                unnormalized_probs.append(G[dst][dst_nbr]['weight'])
            else:
                unnormalized_probs.append(G[dst][dst_nbr]['weight'] / q)
        norm_const = sum(unnormalized_probs)
        normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]

        return alias_setup(normalized_probs)

    def preprocess_transition_probs(self):
        """
        Preprocessing of transition probabilities for guiding the random walks.
        """
        G = self.G
        is_directed = self.is_directed

        alias_nodes = {}
        for node in G.nodes():
            unnormalized_probs = [G[node][nbr]['weight'] for nbr in sorted(G.neighbors(node))]
            norm_const = sum(unnormalized_probs)
            normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]
            alias_nodes[node] = alias_setup(normalized_probs)

        alias_edges = {}

        if is_directed:
            for edge in G.edges():
                alias_edges[edge] = self.get_alias_edge(edge[0], edge[1])
        else:
            for edge in G.edges():
                alias_edges[edge] = self.get_alias_edge(edge[0], edge[1])
                alias_edges[(edge[1], edge[0])] = self.get_alias_edge(edge[1], edge[0])

        self.alias_nodes = alias_nodes
        self.alias_edges = alias_edges

        return


def alias_setup(probs):
    """
    Compute utility lists for non-uniform sampling from discrete distributions.
    Refer to https://hips.seas.harvard.edu/blog/2013/03/03/the-alias-method-efficient-sampling-with-many-discrete-outcomes/
    for details
    """
    K = len(probs)
    q = np.zeros(K)
    J = np.zeros(K, dtype=np.int)

    smaller = []
    larger = []
    for kk, prob in enumerate(probs):
        q[kk] = K * prob
        if q[kk] < 1.0:
            smaller.append(kk)
        else:
            larger.append(kk)

    while len(smaller) > 0 and len(larger) > 0:
        small = smaller.pop()
        large = larger.pop()

        J[small] = large
        q[large] = q[large] + q[small] - 1.0
        if q[large] < 1.0:
            smaller.append(large)
        else:
            larger.append(large)

    return J, q


def alias_draw(J, q):
    """
    Draw sample from a non-uniform discrete distribution using alias sampling.
    """
    K = len(J)

    kk = int(np.floor(np.random.rand() * K))
    if np.random.rand() < q[kk]:
        return kk
    else:
        return J[kk]
